import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from tqdm import tqdm


# 定义一个两次卷积操作的模块
class DoubleConv(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(DoubleConv, self).__init__()
        self.double_conv = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        return self.double_conv(x)


# 定义 U-Net 网络结构
class UNet(nn.Module):
    def __init__(self, in_channels=1, out_channels=1, features=[64, 128, 256, 512]):
        super(UNet, self).__init__()
        self.downs = nn.ModuleList()
        self.ups = nn.ModuleList()
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        # 下采样
        for feature in features:
            self.downs.append(DoubleConv(in_channels, feature))
            in_channels = feature
        # 瓶颈
        self.bottleneck = DoubleConv(features[-1], features[-1] * 2)
        # 上采样
        for feature in reversed(features):
            self.ups.append(nn.ConvTranspose2d(feature * 2, feature, kernel_size=2, stride=2))
            self.ups.append(DoubleConv(feature * 2, feature))
        # 最终卷积
        self.final_conv = nn.Conv2d(features[0], out_channels, kernel_size=1)

    def forward(self, x):
        skip_connections = []
        for down in self.downs:
            x = down(x)
            skip_connections.append(x)
            x = self.pool(x)
        x = self.bottleneck(x)
        skip_connections = skip_connections[::-1]
        for idx in range(0, len(self.ups), 2):
            x = self.ups[idx](x)
            skip_connection = skip_connections[idx // 2]
            if x.shape != skip_connection.shape:
                x = F.interpolate(x, size=skip_connection.shape[2:], mode='bilinear', align_corners=True)
            x = torch.cat((skip_connection, x), dim=1)
            x = self.ups[idx + 1](x)
        x = self.final_conv(x)  # 移除 sigmoid，直接输出 logits
        return x


class MedicalImageDataset(Dataset):
    def __init__(self, list_file):
        self.data = []
        self.root_dir = '/'.join(list_file.split('/')[:-1])
        with open(list_file, 'r') as file:
            for line in file:
                image_path, label_path = line.strip().split()
                self.data.append((self.root_dir + '/' + image_path, self.root_dir + '/' + label_path))

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        img_path, label_path = self.data[idx]
        image = np.load(img_path)
        mask = np.load(label_path)
        transform = transforms.Compose([
            transforms.ToTensor(),
        ])
        image = transform(image)
        mask = transform(mask).long()
        return image, mask


def train(model, dataloader, criterion, optimizer, device, num_epochs=31):
    model.train()
    for epoch in range(num_epochs):
        loop = tqdm(dataloader, leave=True, desc=f'Epoch {epoch + 1}/{num_epochs}')
        total_loss = 0
        for images, masks in loop:
            images = images.to(device)
            masks = masks.to(device)
            optimizer.zero_grad()
            outputs = model(images)
            masks = masks.squeeze(1)
            loss = criterion(outputs, masks)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
            loop.set_postfix(loss=loss.item())
        avg_loss = total_loss / len(dataloader)
        print(f'Epoch {epoch + 1}, Average Loss: {avg_loss:.4f}')
        if (epoch + 1) % 3 == 0:
            torch.save(model.state_dict(), f'../pt_file/synapse-unet-{epoch}.pt')


if __name__ == '__main__':
    list_file = '/Volumes/For_Mac/dateset/Synapse_npy/train_list.txt'
    dataset = MedicalImageDataset(list_file)
    dataloader = DataLoader(dataset, batch_size=8, shuffle=True, num_workers=4)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = UNet(in_channels=1, out_channels=9).to(device)
    criterion = nn.CrossEntropyLoss()  # 修改为使用 CrossEntropyLoss
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    train(model, dataloader, criterion, optimizer, device)
